from graphviz import Digraph
from datetime import datetime

# 获取当前时间
now = datetime.now()
now_str = now.strftime("%Y%m%d_%H%M%S")

# 创建一个有向图，注释为英文
dot = Digraph(comment='Workflow for Diabetic Retinopathy Classification', format='png')

# 设置节点和边的全局属性
dot.attr('node', shape='ellipse', fontname='FangSong', fontsize='10', style='filled')
dot.attr('edge', fontname='FangSong', fontsize='10', arrowhead='open')

# 添加节点，并为每个节点设置不同的线性渐变颜色
dot.node('A', '原始3D眼底彩色图像', style='radial', gradient='on', fillcolor='#FFCCCC:#FF6666')
dot.node('B', 'U-Net模型', style='radial', gradient='on', fillcolor='#CCFFCC:#66CC66')
dot.node('C', '分割特征（血管、出血等）', style='radial', gradient='on', fillcolor='#CCCCFF:#6666CC')
dot.node('D', '特征融合', style='radial', gradient='on', fillcolor='#FFCCFF:#FF66CC')
dot.node('E1', 'EfficientNetV2特征提取', style='radial', gradient='on', fillcolor='#FFFFCC:#FFCC66')
dot.node('E2', 'GCA注意力机制', style='radial', gradient='on', fillcolor='#CCFFFF:#66CCFF')
dot.node('F', '分类器', style='radial', gradient='on', fillcolor='#FFCC99:#FF9933')
dot.node('G', '五分类结果', style='radial', gradient='on', fillcolor='#99CCFF:#3366CC')

# 添加边
dot.edge('A', 'B', label='输入')
dot.edge('B', 'C', label='分割')
dot.edge('A', 'D', label='原始图像')
dot.edge('C', 'D', label='分割特征')
dot.edge('D', 'E1', label='融合特征输入')
dot.edge('E1', 'E2', label='特征增强')
dot.edge('E2', 'F', label='注意力增强特征')
dot.edge('F', 'G', label='分类')

# 保存和渲染图像
dot.render('dr_classification_workflow'+now_str, view=True, cleanup=True)